import socketio
import asyncio
import logging
from app import SioApp

from zunclient.common.apiclient.exceptions import NotFound as NotFoundContainer
from app.common.libs.time_jws import TimedJSONWebSignatureSerializer as Serializer

from tools import ops
from tools import aio_redis_pool
from utils.utils import stop_thread
from configs.settings import flask_config

from .worker.container_thread import ContainerAgentThread


sio_app = SioApp()
sio: socketio.AsyncServer = sio_app.engineio_server
logger = logging.getLogger(__name__)


class ContainerWorker(socketio.AsyncNamespace):

    WS_SID_CACHE = "ZUN_WS_THREAD_IDENT"

    current_namespace = "/container_exec"

    daemon_thread = {}

    def on_connect(self, sid, environ, auth=None):
        with sio_app.other_asgi_app.wsgi_application.app_context():
            if auth and auth.get("Authorization"):
                auth_token = auth["Authorization"]
                if auth_token.startswith("Bearer"):
                    _, token = auth_token.split()
                    try:
                        serializer = Serializer(flask_config.secret_key)
                        serializer.loads(token)
                    except Exception as e:
                        logger.error("token解析失败", exc_info=e)
                        asyncio.run(sio.disconnect(sid))
                        return
                    else:
                        return
            else:
                logger.error(f"{environ['HTTP_ORIGIN']} 试图非法连接到容器, auth={auth}")
                asyncio.run(sio.disconnect(sid))

    async def on_disconnect(self, sid, namespace=None):
        h_key = f"{self.WS_SID_CACHE}::{sid}"
        _cache = await aio_redis_pool.execute("HGETALL", h_key)  # noqa
        cache_info = {
            key.decode(): val.decode() for key, val in _cache.items()
        }

        try:
            vm_uuid = cache_info['vm_uuid']
            container_uuid = cache_info['container_uuid']
        except KeyError:
            try:
                await aio_redis_pool.execute("DEL", h_key)
            except Exception as e:
                logger.error(f"链接: {sid} 清理缓存数据失败", exc_info=e)
            return

        stop_flag = f"{sid}:{container_uuid}::{vm_uuid}::STOP"

        thread_id = self.daemon_thread.get(sid, None)
        try:
            await aio_redis_pool.publish(f"container_channel::{sid}::{cache_info['vm_uuid']}", stop_flag.encode())
        except Exception as e:
            logger.error(f"链接: {sid} 清理线程 {thread_id} 失败，开始强制清理", exc_info=e)
            try:
                stop_thread(thread_id)
            except Exception as _:  # noqa
                logger.error(f"链接: {sid} 清理守护线程 {thread_id} 失败", exc_info=e)
        finally:
            try:
                await sio.close_room(container_uuid, namespace=self.current_namespace)
                await aio_redis_pool.execute("DEL", h_key)
            except Exception as e:
                logger.error(f"链接: {sid} 清理缓存数据失败", exc_info=e)
            logger.info(f"{sid} 连接断开")

    async def on_container_setup(self, sid, data):
        """链接容器，初始化ws链接"""

        vm_uuid = data["container_id"]
        from app.models import RangeVm, ActiveRangeVm

        with sio_app.other_asgi_app.wsgi_application.app_context():
            vm = RangeVm.query.filter(RangeVm.vm_uuid == vm_uuid).first()
            if not vm:
                vm = ActiveRangeVm.query.filter(ActiveRangeVm.new_vm_uuid == vm_uuid).first()
            if not vm:
                await self.emit('from_container_data', "虚拟机不存在")
                await sio.disconnect(sid, self.current_namespace)
                return
            elif not vm.op_vm_uuid:
                await self.emit('from_container_data', "虚拟机未开机")
                await sio.disconnect(sid, self.current_namespace)
                return

            container_uuid = vm.op_vm_uuid

        room_name = f"{sid}::{vm_uuid}"
        sio.enter_room(sid, room_name, namespace=self.current_namespace)

        try:
            conn_info = ops.zun_operate.execute(
                container_uuid, "/bin/sh", run=False, interactive=True
            )
        except NotFoundContainer:
            logger.error(f"容器: {vm_uuid} 未在当前可用域中找到")
            await sio.close_room(room_name, namespace=self.current_namespace)
            await sio.disconnect(sid, namespace=self.current_namespace)
            return

        # 调整尺寸
        # ops.zun_operate.execute_resize(
        #     container_uuid, conn_info["exec_id"], data["cols"], data["rows"]
        # )
        tt = ContainerAgentThread(
            container_uuid,
            vm_uuid,
            self.current_namespace,
            sid=sid,
            ws_url=conn_info["proxy_url"],
            queue_url=sio.manager.redis_url,
            # daemon=True
        )
        tt.start()
        self.daemon_thread.update({sid: tt.ident})

        h_key = f"{self.WS_SID_CACHE}::{sid}"

        await aio_redis_pool.execute(
            "HSET", h_key,
            "zun_thread_id", str(tt.ident), "container_uuid", container_uuid, "vm_uuid", vm_uuid
        )
        logger.info(f"容器{container_uuid} 缓存信息存储成功")

    async def on_container_tty(self, sid, data):  # noqa
        """交互操作"""
        input_data = data["input"]
        vm_uuid = data["container_id"]
        await aio_redis_pool.publish(f"container_channel::{sid}::{vm_uuid}", input_data.encode())

    async def on_my_event(self, sid, data):
        logger.error(f"{sid} 接收： {data}")
        await self.emit('my_response', data)


sio.register_namespace(ContainerWorker(ContainerWorker.current_namespace))
